package io.github.eric.urpc.server;

import io.github.eric.urpc.core.conn.Connection;
import io.github.eric.urpc.core.conn.ConnectionListener;
import io.github.eric.urpc.core.conn.ConnectionManager;
import io.github.eric.urpc.core.rpc.msg.ApplicationInfoReq;
import io.github.eric.urpc.core.rpc.msg.ApplicationInfoRes;
import io.github.eric.urpc.core.util.ConstantsUtil;
import io.netty.util.Timeout;
import io.netty.util.TimerTask;

import java.net.InetSocketAddress;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class ServerConnectionManager implements ConnectionManager, ConnectionListener {
	private final ServerContext serverContext;
	private final Map<String, Connection> connectionMap;
	private final Map<String, ConnectionGroup> connectionGroupMap;

	protected ServerConnectionManager(ServerContext serverContext) {
		this.serverContext = serverContext;
		watchConnection();
		connectionMap = new ConcurrentHashMap<>();
		connectionGroupMap = new ConcurrentHashMap<>();
	}

	@Override
	public void add(Connection conn) {
		Connection old = connectionMap.put(conn.getId(), conn);
		if (old == null) {
			serverContext.getLogger().info("[urpc server] on connection added: {}, local addr {}, remote addr {}",
					conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress());
		}
	}

	@Override
	public void remove(Connection conn) {
		Connection old = connectionMap.remove(conn.getId());
		if (old != null) {
			Object attr = old.getAttribute(ConstantsUtil.APPLICATION_NAME);
			if (attr != null) {
				String applicationName = (String) attr;
				ConnectionGroup connectionGroup = connectionGroupMap.get(applicationName);
				if (connectionGroup != null) {
					connectionGroup.removeConnection(old);
				}
			}
			serverContext.getLogger().info("[urpc server] on connection removed: {}, local addr {}, remote addr {}",
					old.getId(), old.getLocalAddress(), old.getRemoteAddress());
		}
	}

	@Override
	public void shutdown() {
		for (Map.Entry<String, Connection> entry : connectionMap.entrySet()) {
			Connection conn = entry.getValue();
			try {
				conn.close();
			} catch (Exception e) {
				serverContext.getLogger().error("[urpc server] close conn {} error, local addr {}, remote addr {}.",
						conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress(), e);
			}
		}
	}

	public Collection<Connection> getConnectionsByApplicationName(String applicationName) {
		ConnectionGroup connectionGroup = connectionGroupMap.get(applicationName);
		if (connectionGroup == null) {
			return Collections.emptySet();
		}
		return connectionGroup.getAllConnections();
	}

	public Collection<Connection> getConnectionsByApplicationNameAndId(String applicationName, String applicationId) {
		ConnectionGroup connectionGroup = connectionGroupMap.get(applicationName);
		if (connectionGroup == null) {
			return Collections.emptySet();
		}
		return connectionGroup.getConnections(applicationId);
	}

	@Override
	public Collection<Connection> getConnections(Predicate<Connection> predicate) {
		if (predicate == null) {
			return getConnections();
		}
		Collection<Connection> connections = getConnections();
		return connections.stream().filter(predicate).collect(Collectors.toList());
	}

	@Override
	public Collection<Connection> getConnections() {
		return Collections.unmodifiableCollection(connectionMap.values());
	}

	@Override
	public void watchConnection() {
		if (serverContext.isShuttingDown()) {
			return;
		}
		try {
			serverContext.getHashedWheelTimer().newTimeout(new TimerTask() {
				public void run(Timeout timeout) throws Exception {
					try {
						int total = 0;
						Map<String, Integer> result = new HashMap<>();
						Collection<Connection> connections = getConnections();
						for (Connection conn : connections) {
							InetSocketAddress remoteAddress = conn.getRemoteAddress();
							if (remoteAddress == null) {
								int cur = result.computeIfAbsent("null", k -> 0);
								result.put("null", ++cur);
							} else {
								String host = remoteAddress.getHostString();
								int cur = result.computeIfAbsent(host, k -> 0);
								result.put(host, ++cur);
							}
							++total;
						}
						StringBuilder sb = new StringBuilder();
						sb.append("conn total: ").append(total);
						sb.append("\r\n");
						for (Map.Entry<String, Integer> entry : result.entrySet()) {
							sb.append("conn group by host '").append(entry.getKey()).append("': ").append(entry.getValue());
						}
						serverContext.getLogger().info("[urpc server] current server conn info\r\n{}", sb);
					} catch (Throwable t) {
						serverContext.getLogger().error("[urpc server] log server conn info error", t);
					} finally {
						watchConnection();
					}
				}
			}, 60, TimeUnit.SECONDS);
		} catch (Exception e) {
			serverContext.getLogger().error("[urpc server] create log server conn task error!", e);
		}
	}

	@Override
	public void onConnected(Connection conn) {
		throw new UnsupportedOperationException();
	}

	@Override
	public void onDisConnected(Connection conn) {
		remove(conn);
	}

	@Override
	public void onAccepted(Connection conn) {
		add(conn);
	}

	@Override
	public ApplicationInfoRes updateConnectionInitStatus(String connId, ApplicationInfoReq applicationInfoReq) {
		ApplicationInfoRes applicationInfoRes = new ApplicationInfoRes();

		Connection conn = connectionMap.get(connId);
		if (conn == null) {
			applicationInfoRes.setError(String.format("[urpc server] update connection init status error, can't find conn %s in server connection map", connId));
			serverContext.getLogger().error("[urpc server] update connection {} init status error, can't find conn in server connection map", connId);
			return applicationInfoRes;
		}
		try {
			String applicationName = applicationInfoReq.getApplicationName();
			conn.setInitStatusCompleted();
			conn.putAttribute(ConstantsUtil.APPLICATION_NAME, applicationName);
			conn.putAttribute(ConstantsUtil.APPLICATION_ID, applicationInfoReq.getApplicationId());
			ConnectionGroup connectionGroup = connectionGroupMap.computeIfAbsent(applicationName, k -> new ConnectionGroup(applicationName));
			connectionGroup.addConnection(conn);
			serverContext.getLogger().info("[urpc server] update connection {} init status complete", connId);
		} catch (Exception e) {
			conn.setInitStatusFailed();
			applicationInfoRes.setError(e.getMessage());
			serverContext.getLogger().error("[urpc server] update conn {} init status error", connId, e);
		}
		return applicationInfoRes;
	}

	@Override
	public void updateConnectionInitStatus(String connId, boolean succeed) {
		throw new UnsupportedOperationException();
	}

	public void setConnectionAttributes(String connId, Map<String, Object> attributes) {
		Connection conn = connectionMap.get(connId);
		if (conn == null || attributes == null || attributes.isEmpty()) {
			return;
		}
		for (Map.Entry<String, Object> entry : attributes.entrySet()) {
			conn.putAttribute(entry.getKey(), entry.getValue());
		}
	}
}
